#!/usr/bin/env python
import os
try:
    from astropy.io import fits as pf
except ImportError:
    import pyfits as pf

def getMCube(srcMaps, expCube, binnedExpMap, srcModel, irfs, outfile):
    cmdline = ['gtmodel', 'irfs='+irfs, 'outfile='+outfile,
               'srcmaps='+srcMaps, 'expcube='+expCube, 'bexpmap='+binnedExpMap, 'srcmdl='+srcModel,
               'outtype=ccube']
    os.system(' '.join(cmdline))

def sumMCube(srcMaps, expCube, binnedExpMap, srcModel, irfs, formatter, outfile):
    outfiles = []
    srcmapfiles = []
    for i in xrange(4):
        folder = formatter % i
        if not folder:
            continue
        elif not folder.endswith('/'):
            folder += '/'

        if os.path.exists(folder):
            print 'Add models in %s into sumLike!' % folder
            srcmaps_file = folder + srcMaps
            binnedExpMap_file = folder + binnedExpMap
            outfile2 = folder + outfile.split('/')[-1]

            getMCube(srcmaps_file, expCube, binnedExpMap_file, srcModel, irfs, outfile2)
            outfiles.append(outfile2)
            srcmapfiles.append(srcmaps_file)

    if len(outfiles) == 0:
        print 'No mcube is available!'
    elif len(outfiles) == 1:
        print 'MCube saves to '+ outfiles[0]
    else:
        for i, out in enumerate(outfiles):
            if i == 0:
                print 'Open mcube ' + out
                hdusave = pf.open(out)
                ccube = pf.open(srcmapfiles[0])[0]
            else:
                print 'Merge ' + out
                hdulist = pf.open(out)
                hdusave[0].data += hdulist[0].data
                hdulist.close()

                hdulist2 = pf.open(srcmapfiles[i])
                ccube.data += hdulist2[0].data
                hdulist2.close()
                del hdulist, hdulist2
        print 'Save to ' + outfile
        try:
            hdusave.writeto(outfile, overwrite=True)
        except:
            hdusave.writeto(outfile, clobber=True)

        try:
            ccube.writeto('_'.join(outfile.split('_')[:-1])+'_ccube_combined.fits', overwrite=True)
        except:
            ccube.writeto('_'.join(outfile.split('_')[:-1])+'_ccube_combined.fits', clobber=True)


def cli():
    import argparse

    parser = argparse.ArgumentParser(description='This is a binned likelihood analysis tools based on Fermi pyLikelihood package')
    parser.add_argument("srcMaps", type=str, help='Source Maps')
    parser.add_argument("expCube", type=str, help='Live Time Cube')
    parser.add_argument("binnedExpMap", type=str, help='Exposure Cube')
    parser.add_argument("srcModel", type=str, help='model.xml')
    parser.add_argument("irfs", type=str, help='IRFs')
    parser.add_argument("fmt", type=str, help='The formatter of the model folders.')
    parser.add_argument("outfile", type=str, help='output mcube file')
    args = parser.parse_args()

    sumMCube(args.srcMaps, args.expCube, args.binnedExpMap, args.srcModel, args.irfs, args.fmt, args.outfile)

if __name__ == '__main__': cli()
